#!/usr/bin/python -S
"""
ui.py
"""

__author__ = 'Andy Chu'


import sys
import time

import tnet


class Error(Exception):
  pass


CURSOR_UP = '\033[A'
ERASE_LINE = '\033[2K'


_LINE = '%4d  -->  %-15s  -->  %4d tasks in %.1fs   (%.1f tasks/s)'
#using %d processes'

class PumpState(object):
  """
  Accumulates state from the status stream of a ParallelPump instance.
  """

  def __init__(self):
    self.initialized = False
    self.max_procs = 0  # emitted once at the beginning
    self.start_time = 0  # ditto

    self.cur_time = 0  # current time (according to this process)

    # tracked by parent: number written, number finished
    self.num_pending = 0  # emitted on every task finished
    self.num_finished = 0  # ditto

    # tracked by parent per child
    self.pending = {}  # pid -> # pending
    self.finished = {}  # pid -> # finished

  def on_value(self, value):
    """
    A value is available from the input stream.  Update our internal state.
    """
    for k, v in value.iteritems():
      if hasattr(self, k):
        setattr(self, k, v)
      else:
        raise RuntimeError('Invalid name %r in record %r' % (k, value))

    # We got a real value, now we're able to show status
    if 'num_finished' in value:
      self.initialized = True

  def _status_line(self, pending, name, finished):
    elapsed = self.cur_time - self.start_time
    throughput = finished / elapsed
    return _LINE % (pending, name, finished, elapsed, throughput)

  def line(self, name=None):
    if not self.initialized:
      return 'UNINIT'
    name = name or 'xmap root'
    status = self._status_line(
        self.num_pending, name, self.num_finished)
    # could display max_procs -- different than actual number
    # self.max_procs
    return status

  def child_lines(self):
    if not self.initialized:
      return []

    # This should always be the case because we initialize them at the same time
    # as the process
    assert len(self.finished) == len(self.pending)
    lines = []
    for pid in self.pending:
      lines.append(
          self._status_line(self.pending[pid], str(pid), self.finished[pid]))

    return lines


class XmapUi(object):
  """
  You can direct a ParallelPump status stream to this object.

  Additionally, if a child of this xmap process is ANOTHER xmap process, you can
  instruct it to multiplex its status/output together on stdout as follows:

    $ xmap --output @1 --status-out @2    (child)

  And then this process should be invoked like:

    $ xmap --proc-set foo --child-demux output,status  (parent)

  XmapTool will direct child status streams to on_child_chunk() and
  on_child_chunk(), with a PID.

  NOTES:

  - This could have some options about verbosity.  e.g.  do we want to display
  pending/finished count for each process on a remote machine, or just a
  summary?

  - What if lines are too long?  Then you will get "crawling" on the
  screen.  We could assume a fixed-width layout and truncate things...

  MOCK:

  1000   ->  xmap     -> 15 tasks finished in 30s  (0.5 tasks/s, 2 processes)
  400    ->    3969   ->  8 tasks finished in 30s (0.3 tasks/s, 2 processes)
  200    ->      3966 ->  8 tasks finished in 30s (0.3 tasks/s, 2 processes)
  200    ->    3968   ->  9 tasks finished in 31s (0.3 tasks/s, 3 processes)

  6 tasks in .1fs seconds using %d processes
    (for each of the 4 processes)
    3 tasks in .1fs seconds using %d processes
    3 tasks in .1fs seconds using %d processes


  ERRORS:    TODO: test these out

  If you output multiplex, and don't set xmap to parse it, then you'll just get
  everything on stdout... that's OK

  If you try to have xmap demux a non-mux'ed child, it won't see an integer
  stream ID.  I suppose it depends on the output of the child?  It's possible it
  could parse.  Maybe muxing needs a better integrity check.

  [ header payload header payload ... ]
  [ ID payload ID payload ... ]

  Maybe the first record should be negotiation.  If --demux=output,status is
  passed, then it should expect a schema to be sent as the first record.

  e.g. { "schema": { stdout:1, stderr:2 } }

  If they don't match, then it's an error.
  """

  def __init__(self, animate=True):
    self.pump_state = PumpState()
    self.children = {}  # PID -> PumpState()
    self.lines_drawn = 0
    self.animate = animate

  def on_child_chunk(self, pid, chunk):
    #print >>sys.stderr, 'CCHUNK', pid, chunk

    # NOTE: The PIDs in the pending dict get turned into strings because TNET
    # requires string keys.
    value = tnet.loads(chunk)
    if pid not in self.children:
      self.children[pid] = PumpState()
    self.children[pid].on_value(value)

    self._redraw()

  def on_child_close(self, pid):
    print >>sys.stderr, 'XmapUi on_child_chunk', pid

  def on_value(self, value):
    """
    A value is available from the input stream.  Update our internal state.
    """
    self.pump_state.on_value(value)

    self._redraw()

  def on_close(self):
    """
    """
    # TODO: Display things one last time here?
    #print >>sys.stderr, 'XmapUi on_close called'

    # TODO: This should move to XmapUi on_close!
    # Extra newline prevents the prompt from appearing on the same line as
    # status
    # Synchronous
    #print >>sys.stderr, ''

  def _erase(self):
    for i in xrange(self.lines_drawn):
      sys.stderr.write(CURSOR_UP + ERASE_LINE)

  def _redraw(self):
    """
    We get events from the ParallelPump every second.
    We MAY also get events from its (demuxed) child streams, if they're xmap
    processes.  These will be generated every second on a remote machine.

    We update our internal state each time one of those happens, and then call
    redraw().  Each time we keep track of the number of lines we drew, and then
    erase them before redrawing.
    """
    if self.animate:
      self._erase()
    f = sys.stderr

    INDENT = '  '

    i = 0
    print >>f, self.pump_state.line()
    i += 1

    if self.children:
      for pid, state in self.children.iteritems():
        #print >>f, 'PID', pid 
        #print >>f, 'state', state.__dict__
        print >>f, INDENT + state.line('xmap(%s)' % pid)
        i += 1

        lines = state.child_lines()
        for line in lines:
          print >>f, INDENT*2 + line
          i += 1
    else:
      # Just print direct children
      lines = self.pump_state.child_lines()
      for line in lines:
        print >>f, INDENT + line
        i += 1

    self.lines_drawn = i
